{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_omgylxzm5i9"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "f0A2utIXbPc5"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eriCSHTznS4U"
      },
      "source": [
        "# Partial Differential Equations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uYCNQT4snWr6"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r1/tutorials/non-ml/pdes.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r1/tutorials/non-ml/pdes.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uYCNQT4snWr6"
      },
      "source": [
        "> Note: This is an archived TF1 notebook. These are configured\n",
        "to run in TF2's \n",
        "[compatbility mode](https://www.tensorflow.org/guide/migrate)\n",
        "but will run in TF1 as well. To use TF1 in Colab, use the\n",
        "[%tensorflow_version 1.x](https://colab.research.google.com/notebooks/tensorflow_version.ipynb)\n",
        "magic."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zxQvbi5gnyMm"
      },
      "source": [
        "TensorFlow isn't just for machine learning. Here you will use TensorFlow to simulate the behavior of a [partial differential equation](https://en.wikipedia.org/wiki/Partial_differential_equation). You'll simulate the surface of a square pond as a few raindrops land on it.\n",
        "\n",
        "## Basic setup\n",
        "\n",
        "A few imports you'll need."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FG6DLet6ol3j"
      },
      "outputs": [],
      "source": [
        "#Import libraries for simulation\n",
        "import tensorflow.compat.v1 as tf\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "#Imports for visualization\n",
        "import PIL.Image\n",
        "from io import BytesIO\n",
        "from IPython.display import clear_output, Image, display\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7vd7rHS0oqEF"
      },
      "source": [
        "A function for displaying the state of the pond's surface as an image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fJ8SpYYUoq6G"
      },
      "outputs": [],
      "source": [
        "def DisplayArray(a, fmt='jpeg', rng=[0,1]):\n",
        "  \"\"\"Display an array as a picture.\"\"\"\n",
        "  a = (a - rng[0])/float(rng[1] - rng[0])*255\n",
        "  a = np.uint8(np.clip(a, 0, 255))\n",
        "  f = BytesIO()\n",
        "  PIL.Image.fromarray(a).save(f, fmt)\n",
        "  clear_output(wait = True)\n",
        "  display(Image(data=f.getvalue()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NjiZ2_6Mou13"
      },
      "source": [
        "Here you start an interactive TensorFlow session for convenience in playing around. A regular session would work as well if you were doing this in an executable .py file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cH82JlsPozdV"
      },
      "outputs": [],
      "source": [
        "sess = tf.InteractiveSession()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hbk97yero5a9"
      },
      "source": [
        "## Computational convenience functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XVomNV1OpBbX"
      },
      "outputs": [],
      "source": [
        "def make_kernel(a):\n",
        "  \"\"\"Transform a 2D array into a convolution kernel\"\"\"\n",
        "  a = np.asarray(a)\n",
        "  a = a.reshape(list(a.shape) + [1,1])\n",
        "  return tf.constant(a, dtype=1)\n",
        "\n",
        "def simple_conv(x, k):\n",
        "  \"\"\"A simplified 2D convolution operation\"\"\"\n",
        "  x = tf.expand_dims(tf.expand_dims(x, 0), -1)\n",
        "  y = tf.nn.depthwise_conv2d(x, k, [1, 1, 1, 1], padding='SAME')\n",
        "  return y[0, :, :, 0]\n",
        "\n",
        "def laplace(x):\n",
        "  \"\"\"Compute the 2D laplacian of an array\"\"\"\n",
        "  laplace_k = make_kernel([[0.5, 1.0, 0.5],\n",
        "                           [1.0, -6., 1.0],\n",
        "                           [0.5, 1.0, 0.5]])\n",
        "  return simple_conv(x, laplace_k)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9gBib2lpINO"
      },
      "source": [
        "## Define the PDE\n",
        "\n",
        "Our pond is a perfect 500 x 500 square, as is the case for most ponds found in nature."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7faiwBQhpK1Z"
      },
      "outputs": [],
      "source": [
        "N = 500"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U_DscmhfpPs0"
      },
      "source": [
        "Here you create a pond and hit it with some rain drops."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mtk8t0IOpSrb"
      },
      "outputs": [],
      "source": [
        "# Initial Conditions -- some rain drops hit a pond\n",
        "\n",
        "# Set everything to zero\n",
        "u_init = np.zeros([N, N], dtype=np.float32)\n",
        "ut_init = np.zeros([N, N], dtype=np.float32)\n",
        "\n",
        "# Some rain drops hit a pond at random points\n",
        "for n in range(40):\n",
        "  a,b = np.random.randint(0, N, 2)\n",
        "  u_init[a,b] = np.random.uniform()\n",
        "\n",
        "DisplayArray(u_init, rng=[-0.1, 0.1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5vzdx9rHpXsl"
      },
      "source": [
        "Now you specify the details of the differential equation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c6uj8LFDpaZO"
      },
      "outputs": [],
      "source": [
        "# Parameters:\n",
        "# eps -- time resolution\n",
        "# damping -- wave damping\n",
        "eps = tf.placeholder(tf.float32, shape=())\n",
        "damping = tf.placeholder(tf.float32, shape=())\n",
        "\n",
        "# Create variables for simulation state\n",
        "U  = tf.Variable(u_init)\n",
        "Ut = tf.Variable(ut_init)\n",
        "\n",
        "# Discretized PDE update rules\n",
        "U_ = U + eps * Ut\n",
        "Ut_ = Ut + eps * (laplace(U) - damping * Ut)\n",
        "\n",
        "# Operation to update the state\n",
        "step = tf.group(\n",
        "  U.assign(U_),\n",
        "  Ut.assign(Ut_))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eAjwNRjTppN-"
      },
      "source": [
        "## Run the simulation\n",
        "\n",
        "This is where it gets fun -- running time forward with a simple for loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jJLvEydzprsy"
      },
      "outputs": [],
      "source": [
        "# Initialize state to initial conditions\n",
        "tf.global_variables_initializer().run()\n",
        "\n",
        "# Run 1000 steps of PDE\n",
        "for i in range(1000):\n",
        "  # Step simulation\n",
        "  step.run({eps: 0.03, damping: 0.04})\n",
        "\n",
        "# Show final image\n",
        "DisplayArray(U.eval(), rng=[-0.1, 0.1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8AcEDQfbpyDT"
      },
      "source": [
        "Look! Ripples!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "pdes.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
